# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: leon

import os
import datetime

from absl import app, flags


def iterate_date(start_date: datetime.date, end_date: datetime.date, reverse=False):
  cur_date = start_date
  if reverse:
    while cur_date > end_date:
      yield cur_date
      cur_date -= datetime.timedelta(days=1)
  else:
    while cur_date < end_date:
      yield cur_date
      cur_date += datetime.timedelta(days=1)


def iterate_and_create_symlink(src_root, dst_root, trading_dates):
  machines = [
    "feed-01.ap-northeast-2.aws",
    "feed-04.ap-northeast-1.aws",
    "feed-02.ap-northeast-2.aws",
    "feed-05.ap-northeast-1.aws",
    "feed-05.cn-hongkong.aliyun",
    "feed-02.eu-west-1.aws",
    "feed-02.us-east-1.aws",
    "feed-01.ap-northeast-1.aws.huobi",
  ]
  for machine in machines:
    for trading_date in trading_dates:
      date_src_dir = os.path.join(src_root, machine, trading_date.strftime('%Y%m%d'))
      date_dst_dir = os.path.join(dst_root, machine, trading_date.strftime('%Y%m%d'))
      if os.path.exists(date_src_dir):
        os.symlink(date_src_dir, date_dst_dir)


def main(argv):
  src_root = flags.FLAGS.src_root
  dst_root = flags.FLAGS.dst_root
  start_date = flags.FLAGS.start_date
  end_date = flags.FLAGS.end_date

  start_date = datetime.datetime.strptime(start_date, '%Y%m%d').date()
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d').date()
  trading_dates = [td for td in iterate_date(start_date, end_date)]

  iterate_and_create_symlink(src_root, dst_root, trading_dates)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('src_root', None, 'Source fastfeed root dir in archive NFS.')
  flags.DEFINE_string('dst_root', None, 'Destination fastfeed root dir in coin-2 NFS.')
  flags.DEFINE_string('start_date', None, 'Start from YYYYmmdd.')
  flags.DEFINE_string('end_date', None, 'End to YYYYmmdd, not included.')

  app.run(main)
